﻿#include "nginx_memory_pool.h"
#include <cstdio>
#include <cstdlib>

void* MemoryPool::ngx_alloc(size_t size) {
    void* p;

    p = malloc(size);
    if (p == nullptr) {
        fprintf(stderr, "malloc(%I64u) failed\n", size);
    }

    return p;
}

void* MemoryPool::ngx_memalign(ngx_uint_t alignment, size_t size) {
    return ngx_alloc(size);
}


bool MemoryPool::ngx_create_pool(size_t size) {
    ngx_pool_t* p;

    p = (ngx_pool_t*)ngx_memalign(NGX_POOL_ALIGNMENT, size);
    if (p == nullptr) {
        pool_ = p;
        return false;
    }

    p->d.last = (u_char*)p + sizeof(ngx_pool_t);
    p->d.end = (u_char*)p + size;
    p->d.next = nullptr;
    p->d.failed = 0;

    size = size - sizeof(ngx_pool_t);
    p->max = (size < NGX_MAX_ALLOC_FROM_POOL) ? size : NGX_MAX_ALLOC_FROM_POOL;

    p->current = p;
    p->large = nullptr;
    p->cleanup = nullptr;

    pool_ = p;

    return true;
}

void MemoryPool::ngx_free(void* data) {
    free(data);
}

void MemoryPool::ngx_destroy_pool() {
    ngx_pool_t* p, * n;
    ngx_pool_large_t* l;
    ngx_pool_cleanup_t* c;

    for (c = pool_->cleanup; c; c = c->next) {
        if (c->handler) {
            fprintf(stderr, "run cleanup: %p\n", c);
            c->handler(c->data);
        }
    }

    for (l = pool_->large; l; l = l->next) {
        if (l->alloc) {
            ngx_free(l->alloc);
        }
    }

    for (p = pool_, n = pool_->d.next; /* void */; p = n, n = n->d.next) {
        ngx_free(p);

        if (n == nullptr) {
            break;
        }
    }
}

void* MemoryPool::ngx_palloc_small(size_t size, bool align) {
    u_char* m;
    ngx_pool_t* p;

    p = pool_->current;

    do {
        m = p->d.last;

        if (align) {
            m = ngx_align_ptr(m, NGX_ALIGNMENT);
        }

        if ((size_t)(p->d.end - m) >= size) {
            p->d.last = m + size;

            return m;
        }

        p = p->d.next;

    } while (p);

    return ngx_palloc_block(size);
}
void* MemoryPool::ngx_palloc_block(size_t size) {
    u_char* m;
    size_t       psize;
    ngx_pool_t* p, * new_mem;

    psize = (size_t)(pool_->d.end - (u_char*)pool_);

    m = (u_char*)ngx_memalign(NGX_POOL_ALIGNMENT, psize);
    if (m == nullptr) {
        return nullptr;
    }

    new_mem = (ngx_pool_t*)m;

    new_mem->d.end = m + psize;
    new_mem->d.next = nullptr;
    new_mem->d.failed = 0;

    m += sizeof(ngx_pool_data_t);
    m = ngx_align_ptr(m, NGX_ALIGNMENT);
    new_mem->d.last = m + size;

    for (p = pool_->current; p->d.next; p = p->d.next) {
        if (p->d.failed++ > 4) {
            pool_->current = p->d.next;
        }
    }

    p->d.next = new_mem;

    return m;
}

bool MemoryPool::ngx_pfree(void* p) {
    ngx_pool_large_t* l;

    for (l = pool_->large; l; l = l->next) {
        if (p == l->alloc) {
            fprintf(stderr, "free: %p\n", l->alloc);
            ngx_free(l->alloc);
            l->alloc = nullptr;

            return true;
        }
    }

    return false;
}

void MemoryPool::ngx_reset_pool() {
    ngx_pool_t* p;
    ngx_pool_large_t* l;

    for (l = pool_->large; l; l = l->next) {
        if (l->alloc) {
            ngx_free(l->alloc);
        }
    }

    // 第一块进行特殊处理
    p = pool_;
    p->d.last = (u_char*)p + sizeof(ngx_pool_t);
    p->d.failed = 0;

    // 从第二块开始
    for (p = p->d.next; p; p = p->d.next) {
        p->d.next = (ngx_pool_t*)((u_char*)p + sizeof(ngx_pool_data_t));
        p->d.failed = 0;
    }

    pool_->current = pool_;
    pool_->large = nullptr;
}
void* MemoryPool::ngx_palloc(size_t size) {
    if (size <= pool_->max) {
        return ngx_palloc_small(size, true);
    }

    return ngx_palloc_large(size);
}

void* MemoryPool::ngx_pnalloc(size_t size) {
    if (size <= pool_->max) {
        return ngx_palloc_small(size, false);
    }
    return ngx_palloc_large(size);
}

void* MemoryPool::ngx_pcalloc(size_t size) {
    void* p;

    p = ngx_palloc(size);
    if (p) {
        ngx_memzero(p, size);
    }

    return p;
}

ngx_pool_cleanup_t* MemoryPool::ngx_pool_cleanup_add(size_t size) {
    ngx_pool_cleanup_t* c;

    // 注意，这里有一个究极大bug。那就是，这里是用malloc来申请ngx_pool_cleanup_t的空间的
    // malloc在申请空间时 不会调用构造函数
    // 会导致ngx_pool_cleanup_t里面的包装器handler包装了野指针
    // 而handler在赋值的时候会尝试释放原本的函数并置空
    // 此时就会导致野指针问题！！！
    // 我决定，不用包装器了
    c = (ngx_pool_cleanup_t*)ngx_palloc(sizeof(ngx_pool_cleanup_t));
    if (c == nullptr) {
        return nullptr;
    }

    if (size) {
        c->data = ngx_palloc(size);
        if (c->data == nullptr) {
            return nullptr;
        }

    }
    else {
        c->data = nullptr;
    }

    c->handler = nullptr;
    c->next = pool_->cleanup;

    pool_->cleanup = c;

    fprintf(stderr, "add cleanup: %p\n", c);

    return c;
}

void* MemoryPool::ngx_palloc_large(size_t size) {
    void* p;
    ngx_uint_t         n;
    ngx_pool_large_t* large;

    p = ngx_alloc(size);
    if (p == nullptr) {
        return nullptr;
    }

    n = 0;

    for (large = pool_->large; large; large = large->next) {
        if (large->alloc == nullptr) {
            large->alloc = p;
            return p;
        }

        if (n++ > 3) {
            break;
        }
    }

    large = (ngx_pool_large_t*)(ngx_palloc_small(sizeof(ngx_pool_large_t), 1));
    if (large == nullptr) {
        ngx_free(p);
        return nullptr;
    }

    large->alloc = p;
    large->next = pool_->large;
    pool_->large = large;

    return p;
}

void MemoryPool::ngx_memzero(void* ptr, size_t size) {
    memset(ptr, 0, size);
}

MemoryPool::MemoryPool(const size_t size) {
    if (!ngx_create_pool(size)) {
        fprintf(stderr, "create memory pool fail!\n");
        exit(EXIT_FAILURE);
    }
}
MemoryPool::~MemoryPool() {
    ngx_destroy_pool();
}